package protocol

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"gitee.com/yuebanlaosiji/gocelery/pkg/types"
	"github.com/google/uuid"
)

var (
	// ErrInvalidMessage is returned when a message is invalid
	ErrInvalidMessage = errors.New("invalid message format")
)

// CeleryProtocol implements the Protocol interface for Celery's wire protocol
type CeleryProtocol struct{}

// NewCeleryProtocol creates a new Celery protocol handler
func NewCeleryProtocol() *CeleryProtocol {
	return &CeleryProtocol{}
}

// CeleryTask represents a Celery task
type CeleryTask struct {
	ID      string                 `json:"id"`
	Task    string                 `json:"task"`
	Args    []interface{}          `json:"args"`
	Kwargs  map[string]interface{} `json:"kwargs"`
	Retries int                    `json:"retries"`
	ETA     string                 `json:"eta,omitempty"`
	Expires string                 `json:"expires,omitempty"`
}

// CeleryResultMessage represents a Celery task result message format
type CeleryResultMessage struct {
	ID        string      `json:"task_id"`
	Status    string      `json:"status"`
	Result    interface{} `json:"result,omitempty"`
	Traceback string      `json:"traceback,omitempty"`
	Error     string      `json:"error,omitempty"`
}

// CreateTaskMessage creates a Celery task message
func (cp *CeleryProtocol) CreateTaskMessage(task types.Task) ([]byte, error) {
	if task.ID == "" {
		task.ID = uuid.New().String()
	}

	// 确保kwargs不为null，Python Celery期望它是一个映射
	if task.Kwargs == nil {
		task.Kwargs = map[string]interface{}{}
	}

	// 构造Celery所需的body格式：[args, kwargs, options]
	bodyArray := []interface{}{
		task.Args,
		task.Kwargs,
		map[string]interface{}{
			"callbacks": nil,
			"errbacks":  nil,
			"chain":     nil,
			"chord":     nil,
		},
	}

	// 将body序列化为JSON
	bodyJSON, err := json.Marshal(bodyArray)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal body: %w", err)
	}

	// 将body编码为base64
	encodedBody := base64.StdEncoding.EncodeToString(bodyJSON)

	// 创建headers
	headers := map[string]interface{}{
		"lang":          "go",
		"task":          task.Name,
		"id":            task.ID,
		"root_id":       task.ID,
		"parent_id":     nil,
		"retries":       task.Retries,
		"eta":           nil,
		"expires":       nil,
		"group":         nil,
		"timelimit":     []interface{}{nil, nil},
		"argsrepr":      fmt.Sprintf("%v", task.Args),
		"kwargsrepr":    fmt.Sprintf("%v", task.Kwargs),
		"origin":        "go-celery",
		"ignore_result": false,
	}

	// 处理时间字段
	if task.ETA != nil {
		headers["eta"] = task.ETA.Format(time.RFC3339)
	}

	if task.ExpiresAt != nil {
		headers["expires"] = task.ExpiresAt.Format(time.RFC3339)
	}

	// 创建完整的消息，包括properties字段
	message := map[string]interface{}{
		"body":             encodedBody,
		"content-encoding": "utf-8",
		"content-type":     "application/json",
		"headers":          headers,
		"properties": map[string]interface{}{
			"correlation_id": task.ID,
			"reply_to":       "",
			"delivery_mode":  2,
			"delivery_info": map[string]interface{}{
				"exchange":    "",
				"routing_key": "celery",
			},
			"priority":      0,
			"body_encoding": "base64",
			"delivery_tag":  uuid.New().String(),
		},
	}

	// 将完整消息序列化为JSON
	return json.Marshal(message)
}

// ParseTaskMessage parses a Celery task message
func (cp *CeleryProtocol) ParseTaskMessage(data []byte) (types.Task, error) {
	// 首先尝试解析为复杂的消息格式（带有headers和properties）
	var message map[string]interface{}
	if err := json.Unmarshal(data, &message); err == nil {
		// 检查是否包含body字段，这表明它是复杂格式
		if _, hasBody := message["body"]; hasBody {
			return cp.parseComplexMessage(data)
		}
	}

	// 如果不是复杂格式，尝试直接解析为简单的任务格式
	var celeryTask CeleryTask
	if err := json.Unmarshal(data, &celeryTask); err != nil {
		return types.Task{}, fmt.Errorf("failed to parse task message: %w", err)
	}

	// 确保kwargs不为null
	if celeryTask.Kwargs == nil {
		celeryTask.Kwargs = map[string]interface{}{}
	}

	// 转换为Task对象
	task := types.Task{
		ID:      celeryTask.ID,
		Name:    celeryTask.Task,
		Args:    celeryTask.Args,
		Kwargs:  celeryTask.Kwargs,
		Retries: celeryTask.Retries,
	}

	// 解析时间字段
	if celeryTask.ETA != "" {
		eta, err := time.Parse(time.RFC3339, celeryTask.ETA)
		if err != nil {
			return types.Task{}, fmt.Errorf("invalid ETA format: %w", err)
		}
		task.ETA = &eta
	}

	if celeryTask.Expires != "" {
		expires, err := time.Parse(time.RFC3339, celeryTask.Expires)
		if err != nil {
			return types.Task{}, fmt.Errorf("invalid Expires format: %w", err)
		}
		task.ExpiresAt = &expires
	}

	return task, nil
}

// parseComplexMessage 解析复杂的Celery/Kombu消息格式
func (cp *CeleryProtocol) parseComplexMessage(data []byte) (types.Task, error) {
	// 解析为映射
	var message map[string]interface{}
	if err := json.Unmarshal(data, &message); err != nil {
		return types.Task{}, fmt.Errorf("failed to parse message: %w", err)
	}

	// 检查是否包含body字段
	bodyRaw, hasBody := message["body"]
	if !hasBody {
		return types.Task{}, errors.New("message does not contain body")
	}

	// 检查headers字段
	headers, hasHeaders := message["headers"].(map[string]interface{})
	if !hasHeaders {
		return types.Task{}, errors.New("message does not contain headers")
	}

	// 从headers中提取任务信息
	taskName, hasTask := headers["task"].(string)
	if !hasTask {
		return types.Task{}, errors.New("headers does not contain task name")
	}

	taskID, hasID := headers["id"].(string)
	if !hasID {
		return types.Task{}, errors.New("headers does not contain task ID")
	}

	// 检查body编码
	var bodyBytes []byte
	var err error

	// 如果body是字符串，检查是否需要解码
	if bodyStr, ok := bodyRaw.(string); ok {
		// 优先检查content-encoding和properties.body_encoding
		bodyEncoding := ""

		// 从properties中获取body_encoding
		if props, ok := message["properties"].(map[string]interface{}); ok {
			if enc, ok := props["body_encoding"].(string); ok {
				bodyEncoding = enc
			}
		}

		if bodyEncoding == "base64" {
			bodyBytes, err = base64.StdEncoding.DecodeString(bodyStr)
			if err != nil {
				return types.Task{}, fmt.Errorf("failed to decode base64 body: %w", err)
			}
		} else {
			bodyBytes = []byte(bodyStr)
		}
	} else {
		// 如果body不是字符串，尝试直接转换
		bodyJSON, err := json.Marshal(bodyRaw)
		if err != nil {
			return types.Task{}, fmt.Errorf("failed to marshal body: %w", err)
		}
		bodyBytes = bodyJSON
	}

	// 尝试解析body
	var argsArray []interface{}
	if err := json.Unmarshal(bodyBytes, &argsArray); err == nil && len(argsArray) >= 2 {
		// Python Celery消息格式：[args, kwargs, options]
		args, _ := argsArray[0].([]interface{})

		// 获取kwargs，确保始终是一个映射
		var kwargs map[string]interface{}
		if kw, ok := argsArray[1].(map[string]interface{}); ok {
			kwargs = kw
		} else {
			kwargs = map[string]interface{}{}
		}

		// 创建任务对象
		task := types.Task{
			ID:     taskID,
			Name:   taskName,
			Args:   args,
			Kwargs: kwargs,
		}

		// 从headers提取其他字段
		if retries, ok := headers["retries"].(float64); ok {
			task.Retries = int(retries)
		}

		return task, nil
	}

	// 如果不是数组格式，尝试解析为任务对象
	var celeryTask CeleryTask
	if err := json.Unmarshal(bodyBytes, &celeryTask); err != nil {
		return types.Task{}, fmt.Errorf("failed to parse task body: %w", err)
	}

	// 确保kwargs不为null
	if celeryTask.Kwargs == nil {
		celeryTask.Kwargs = map[string]interface{}{}
	}

	// 创建任务对象
	task := types.Task{
		ID:      taskID,
		Name:    taskName,
		Args:    celeryTask.Args,
		Kwargs:  celeryTask.Kwargs,
		Retries: celeryTask.Retries,
	}

	// 处理时间字段
	if celeryTask.ETA != "" {
		eta, err := time.Parse(time.RFC3339, celeryTask.ETA)
		if err == nil {
			task.ETA = &eta
		}
	}

	if celeryTask.Expires != "" {
		expires, err := time.Parse(time.RFC3339, celeryTask.Expires)
		if err == nil {
			task.ExpiresAt = &expires
		}
	}

	return task, nil
}

// CreateResultMessage creates a Celery result message
func (cp *CeleryProtocol) CreateResultMessage(result types.TaskResult) ([]byte, error) {
	msg := CeleryResultMessage{
		ID:        result.TaskID,
		Status:    result.Status,
		Result:    result.Result,
		Traceback: result.Traceback,
		Error:     result.Error,
	}

	return json.Marshal(msg)
}

// ParseResultMessage parses a Celery result message
func (cp *CeleryProtocol) ParseResultMessage(data []byte) (types.TaskResult, error) {
	// 尝试解析为Python Celery结果格式
	var pythonResult map[string]interface{}
	if err := json.Unmarshal(data, &pythonResult); err == nil {
		// 检查是否包含标准的Python Celery结果字段
		if status, ok := pythonResult["status"].(string); ok {
			result := types.TaskResult{
				Status: status,
				Result: pythonResult["result"],
			}

			// 获取任务ID（如果存在）
			if taskID, ok := pythonResult["task_id"].(string); ok {
				result.TaskID = taskID
			}

			// 获取traceback（如果存在）
			if traceback, ok := pythonResult["traceback"].(string); ok {
				result.Traceback = traceback
			}

			// 获取错误信息（如果存在）
			if errorMsg, ok := pythonResult["error"].(string); ok {
				result.Error = errorMsg
			}

			return result, nil
		}
	}

	// 如果不是Python格式，尝试解析为标准CeleryResultMessage格式
	var msg CeleryResultMessage
	if err := json.Unmarshal(data, &msg); err != nil {
		return types.TaskResult{}, fmt.Errorf("failed to parse result message: %w", err)
	}

	result := types.TaskResult{
		TaskID:    msg.ID,
		Status:    msg.Status,
		Result:    msg.Result,
		Traceback: msg.Traceback,
		Error:     msg.Error,
	}

	return result, nil
}
